#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：tools 
@File ：共享内存管理.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2024/11/27 下午12:18 
"""

import multiprocessing
import pickle
import pandas as pd
from multiprocessing import shared_memory


class SharedMemoryManager:
	def __init__(self):
		self.shm = None

	def create_shared_memory(self, name, size):
		"""创建共享内存"""
		self.shm = shared_memory.SharedMemory(name=name, create=True, size=size)
		print(f"Shared memory created with name {name} and size {size} bytes")
		return self.shm.name

	def write_data(self, name, value):
		"""将单一变量写入共享内存"""
		self.shm = shared_memory.SharedMemory(name=name, create=False)
		self.shm.buf[:value.nbytes] = value.tobytes()

	def read_data(self, name, dtype):
		"""从共享内存读取单一变量"""
		self.shm = shared_memory.SharedMemory(name=name)
		value = dtype.from_address(self.shm.buf.__address__)
		return value

	def write_structured_data(self, name, data, data_type='pickle'):
		"""将结构化数据写入共享内存，支持pickle和pandas"""
		self.shm = shared_memory.SharedMemory(name=name, create=True, size=1024 * 1024)  # 假设1MB足够
		if data_type == 'pickle':
			serialized_data = pickle.dumps(data)
		elif data_type == 'pandas':
			serialized_data = data.to_bytes(index=False)
		else:
			raise ValueError("Unsupported data type. Use 'pickle' or 'pandas'.")

		self.shm.buf[:len(serialized_data)] = serialized_data

	def read_structured_data(self, name, data_type='pickle'):
		"""从共享内存读取结构化数据，支持pickle和pandas"""
		self.shm = shared_memory.SharedMemory(name=name)
		serialized_data = self.shm.buf[:]

		if data_type == 'pickle':
			data = pickle.loads(serialized_data)
		elif data_type == 'pandas':
			data = pd.from_bytes(serialized_data)
		else:
			raise ValueError("Unsupported data type. Use 'pickle' or 'pandas'.")

		return data

	def delete_shared_memory(self, name):
		"""删除共享内存"""
		self.shm = shared_memory.SharedMemory(name=name)
		self.shm.close()
		self.shm.unlink()
		print(f"Shared memory {name} deleted")


# 使用示例
if __name__ == "__main__":
	manager = SharedMemoryManager()
	shm_name = manager.create_shared_memory('my_shm', 1024)  # 创建1024字节的共享内存

	# 写入单一变量
	value = 42
	manager.write_data(shm_name, value)
	print(f"Value written to shared memory: {manager.read_data(shm_name, int)}")

	# 写入字典数据
	structured_data = {'key': 'value', 'number': 42}
	manager.write_structured_data('my_struct_shm', structured_data, data_type='pickle')

	# 读取字典数据
	data = manager.read_structured_data('my_struct_shm', data_type='pickle')
	print(f"Data read from shared memory: {data}")

	# 写入DataFrame数据
	df = pd.DataFrame({'A': [1, 2], 'B': [3, 4]})
	manager.write_structured_data('my_df_shm', df, data_type='pandas')

	# 读取DataFrame数据
	data = manager.read_structured_data('my_df_shm', data_type='pandas')
	print(data)

	# 删除共享内存
	manager.delete_shared_memory(shm_name)
	manager.delete_shared_memory('my_struct_shm')
	manager.delete_shared_memory('my_df_shm')
